import rospy
from geometry_msgs.msg._PoseWithCovarianceStamped import PoseWithCovarianceStamped
import simpleguitk as simplegui
from geometry_msgs.msg import PoseWithCovarianceStamped,TwistStamped
from sensor_msgs.msg import NavSatFix
from gazebo_msgs.msg import ModelStates
from tf.transformations import euler_from_quaternion
import numpy
import matplotlib
import pylab as plt

WIDTH = 500
LENGTH = 500
name = 'drivingVic'
init_X_Y = [0,0,0]
real_x_log = []
x_log = []
real_y_log = []
y_log = []
speed_x_log = []
speed_y_log = []
real_speed_x_log = []
real_speed_y_log = []
plot_on = 1

class location:
    def __init__(self):
        self._x = 0.0
        self._y = 0.0
        self._z = 0.0
        self._yaw = 0.0
        self._speed_x = 0.0
        self._speed_y = 0.0
        self._acc_x = 0.0
        self._acc_y = 0.0
        self._gy = 0.0

class ekfDataLog:
    def __init__(self):
        rospy.init_node('graph_tracker')
        self._real_loc = location()
        self._estimator = location()
        self._loc_sub = rospy.Subscriber('/LOC/Pose',PoseWithCovarianceStamped, self.estimatedLocation)
        self._speed_sub = rospy.Subscriber('/LOC/Velocity',TwistStamped,self.estimatedSpeed)
        self._real_loc_sub = rospy.Subscriber('/gazebo/model_states',ModelStates, self.realLocation)
        self.init_xyz = [0,0]

    def realLocation(self,obj):
        i = 0
        for objecti in obj.name:
            if objecti == name:
                val = i
                break
            i += 1
        theta = euler_from_quaternion( numpy.array((obj.pose[val].orientation.x, obj.pose[val].orientation.y, obj.pose[val].orientation.z, obj.pose[val].orientation.w), dtype=numpy.float64))
        self._real_loc._x = obj.pose[val].position.x - init_X_Y[0]
        self._real_loc._y = obj.pose[val].position.y - init_X_Y[1]
        self._real_loc._z = obj.pose[val].position.z
        self._real_loc._yaw = theta[2]
        self._real_loc._speed_x = obj.twist[val].linear.x
        self._real_loc._speed_y = obj.twist[val].linear.y
        self._real_loc._gy = obj.twist[val].angular.z
        self.init_xyz = [obj.pose[val].position.x,obj.pose[val].position.y]

    def estimatedLocation(self,obj):
        self._estimator._x = obj.pose.pose.position.x
        self._estimator._y = obj.pose.pose.position.y
        self._estimator._z = obj.pose.pose.position.z
        theta = euler_from_quaternion( numpy.array((obj.pose.pose.orientation.x, obj.pose.pose.orientation.y, obj.pose.pose.orientation.z, obj.pose.pose.orientation.w), dtype=numpy.float64))
        self._estimator._yaw = theta[2]

    def estimatedSpeed(self,obj):
        self._estimator._speed_x = obj.twist.linear.x
        self._estimator._speed_y = obj.twist.linear.y

    def error_calculator(self):
        e_x = self._estimator._x - self._real_loc._x
        e_y = self._estimator._y - self._real_loc._y                
        e_z = self._estimator._z - self._real_loc._z
        e_yaw = self._estimator._yaw - self._real_loc._yaw
        e_s_x = self._estimator._speed_x - self._real_loc._speed_x
        e_s_y = self._estimator._speed_y - self._real_loc._speed_y
        return e_x,e_y,e_z,e_yaw,e_s_x,e_s_y

num = 50
ros_log = ekfDataLog()

while(1):

        real_x_log.append(ros_log._real_loc._x)
        x_log.append(ros_log._estimator._x)
        real_y_log.append(ros_log._real_loc._y)
        y_log.append(ros_log._estimator._y)
        speed_x_log.append(ros_log._estimator._speed_x)
        speed_y_log.append(ros_log._estimator._speed_y)
        real_speed_x_log.append(ros_log._real_loc._speed_x)
        real_speed_y_log.append(ros_log._real_loc._speed_y)

        if len(real_x_log) > num:
            real_x_log.pop(0)
        if len(real_y_log) > num:
            real_y_log.pop(0)
        if len(x_log) > num:
            x_log.pop(0)
        if len(y_log) > num:
            y_log.pop(0)
        if len(real_speed_x_log) > num:
            real_speed_x_log.pop(0)
        if len(speed_x_log) > num:
            speed_x_log.pop(0)
        if len(real_speed_y_log) > num:
            real_speed_y_log.pop(0)
        if len(speed_y_log) > num:
            speed_y_log.pop(0)

        plt.figure(1)
        plt.ion()
        plt.clf()
        plt.subplot(2,1,1)
        plt.plot(real_x_log)
        plt.plot(x_log)
        plt.title("x")
        plt.subplot(2,1,2)
        plt.plot(real_y_log)
        plt.plot(y_log)
        plt.title("y")

        plt.draw()
        plt.ioff()
        
        plt.figure(2)
        plt.ion()
        plt.clf()
        plt.subplot(2,1,1)
        plt.plot(real_speed_x_log)
        plt.plot(speed_x_log)
        plt.title("v_x")
        plt.subplot(2,1,2)
        plt.plot(real_speed_y_log)
        plt.plot(speed_y_log)
        plt.title("v_y")

        plt.draw()
        plt.ioff()
        
        
